package com.mrx.agent;

import com.mrx.agent.plugin.BasePlugin;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.ClassFileLocator;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.dynamic.loading.ClassInjector;
import net.bytebuddy.matcher.ElementMatchers;
import net.bytebuddy.utility.JavaModule;

import java.io.File;
import java.lang.instrument.Instrumentation;
import java.lang.management.ManagementFactory;
import java.util.*;
import java.util.function.Function;
import java.util.jar.JarFile;
import java.util.stream.Collectors;

import static net.bytebuddy.matcher.ElementMatchers.named;

/**
 * @author Mr.X
 * @since 2024-11-04 12:45
 **/
public class XAgent {

    private static final String javaAgent = "-javaagent:";

    public static void premain(final String agentArgs, final Instrumentation inst) throws Exception {
        log("premain: " + agentArgs);
        // -javaagent:D:\agent-1.0.0-jar-with-dependencies.jar=-Duser.language=CN,-Duser.region=CN
        Set<String> args = Collections.emptySet();
        if (agentArgs != null && !agentArgs.isBlank()) {
            args = java.util.Set.of(agentArgs.split(","));
        }
        log("args: " + args);
        // 注入 advice 类
        List<AgentBuilder.Transformer> transformers = injectClass(args);
        if (transformers.isEmpty()) {
            log("transformers is empty");
            return;
        }
        AgentBuilder.Identified.Narrowable builder = new AgentBuilder.Default()
                .with(new ByteBuddyListener())
                .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
                .with(AgentBuilder.InitializationStrategy.NoOp.INSTANCE)
                .with(AgentBuilder.TypeStrategy.Default.REBASE)
                .disableClassFormatChanges()
                .ignore(ElementMatchers.nameStartsWith("net.bytebuddy."))
                .type(named("com.mrx.ExampleMain").or(named("java.lang.ProcessBuilder")));
        // 安装 transformer
        AgentBuilder.Identified.Extendable transform = builder.transform(transformers.get(0));
        for (int i = 1; i < transformers.size(); i++) {
            transform = transform.transform(transformers.get(i));
        }
        transform.installOn(inst);
    }

    private static List<AgentBuilder.Transformer> injectClass(Set<String> args) throws Exception {
        String agentJar = ManagementFactory.getRuntimeMXBean().getInputArguments().stream()
                .filter(it -> it.startsWith(javaAgent)).map(it -> it.replace(javaAgent, ""))
                .map(it -> it.substring(0, it.indexOf(".jar") + 4))
                .filter(it -> it.contains("agent-1.0.0")).findFirst().orElseThrow();
        ClassFileLocator.ForJarFile agentJarFile = new ClassFileLocator.ForJarFile(new JarFile(agentJar));
        // 加载插件类
        List<BasePlugin> plugins = ServiceLoader.load(BasePlugin.class).stream().map(ServiceLoader.Provider::get).toList();
        Map<String, byte[]> resourceMap = plugins.stream()
                .map(BasePlugin::dependencies).flatMap(Collection::stream)
                .distinct().collect(Collectors.toMap(Function.identity(), it -> locateResource(it, agentJarFile)));
        // 注入插件类到 bootstrapClassLoader
        File dir = new File("./test");
        if (!dir.exists() && !dir.mkdirs()) {
            throw new RuntimeException("Unable to create test directory");
        }
        Map<String, Class<?>> clazzMap = ClassInjector.UsingUnsafe.ofBootLoader().injectRaw(resourceMap);
        log("plugins: " + clazzMap.keySet());
        return clazzMap.entrySet().stream().map(it -> {
            BasePlugin supported = plugins.stream().filter(plugin -> plugin.support(it.getKey()))
                    .findFirst().orElse(null);
            if (supported == null) {
                return null;
            }
            supported.setArgs(args, it.getValue());
            return supported.transformers();
        }).filter(Objects::nonNull).flatMap(Collection::stream).toList();
    }

    private static class ByteBuddyListener implements AgentBuilder.Listener {

        @Override
        public void onDiscovery(String s, ClassLoader classLoader, JavaModule javaModule, boolean b) {

        }

        @Override
        public void onTransformation(TypeDescription typeDescription, ClassLoader classLoader, JavaModule javaModule, boolean b, DynamicType dynamicType) {
            log("onTransformation: " + typeDescription.getTypeName());
        }

        @Override
        public void onIgnored(TypeDescription typeDescription, ClassLoader classLoader, JavaModule javaModule, boolean b) {

        }

        @Override
        public void onError(String s, ClassLoader classLoader, JavaModule javaModule, boolean b, Throwable throwable) {
            log("onError: " + s + ", because of: " + throwable.getMessage());
        }

        @Override
        public void onComplete(String s, ClassLoader classLoader, JavaModule javaModule, boolean b) {
            // log("onComplete: " + s);
        }
    }

    private static void log(Object msg) {
        System.out.println(XAgent.class.getSimpleName() + ": " + msg);
    }

    private static byte[] locateResource(String resourceName, ClassFileLocator.ForJarFile agentJarFile) {
        try {
            return agentJarFile.locate(resourceName).resolve();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

}
